package tk.mynoteapp.dev.demo05.routelimit;

import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
import java.util.stream.IntStream;

/**
 * 并发限制了爬虫的抓取能力
 *
 * 查看 PoolingHttpClientConnectionManager 源码，可以注意到有两个重要参数：
 *   - defaultMaxPerRoute=2，也就是同一个主机 / 域名的最大并发请求数为 2。我们的爬虫需要 10 个并发，显然是默认值太小限制了爬虫的效率。
 *   - maxTotal=20，也就是所有主机整体最大并发为 20，这也是 HttpClient 整体的并发度。目前，我们请求数是 10 最大并发是 10，20 不会成
 *     为瓶颈。举一个例子，使用同一个 HttpClient 访问 10 个域名，defaultMaxPerRoute 设置为 10，为确保每一个域名都能达到 10 并发，
 *     需要把 maxTotal 设置为 100。
 */
@RestController
@RequestMapping("routelimit")
public class RouteLimitController {
    private Logger logger = LoggerFactory.getLogger(RouteLimitController.class);

    static CloseableHttpClient httpClient1;
    static CloseableHttpClient httpClient2;

    static {
        httpClient1 = HttpClients.custom().setConnectionManager(new PoolingHttpClientConnectionManager()).build();
        // setMaxConnPerRoute：设置同一域名/主机最大并发请求数10
        // setMaxConnTotal：设置当前主机最大并发数20
        httpClient2 = HttpClients.custom().setMaxConnPerRoute(10).setMaxConnTotal(20).build();
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            try {
                httpClient1.close();
            } catch (IOException e) {
            }
            try {
                httpClient2.close();
            } catch (IOException e) {
            }
        }));
    }

    private int sendRequest(int count, Supplier<CloseableHttpClient> client) throws InterruptedException {
        // 用于计数发送的请求个数
        AtomicInteger atomicInteger = new AtomicInteger();
        // 使用HttpClient从server接口查询数据的任务提交到线程池并行处理
        ExecutorService threadPool = Executors.newCachedThreadPool();
        long begin = System.currentTimeMillis();
        IntStream.rangeClosed(1, count).forEach(i -> {
            threadPool.execute(() -> {
                try (CloseableHttpResponse response = client.get().execute(new HttpGet("http://127.0.0.1:8080/routelimit/server"))) {
                    atomicInteger.addAndGet(Integer.parseInt(EntityUtils.toString(response.getEntity())));
;                } catch (Exception ex) {
                    ex.printStackTrace();
                }
            });
        });
        // 等待count个任务全部执行完毕
        threadPool.shutdown();
        threadPool.awaitTermination(1, TimeUnit.HOURS);
        logger.info("发送 {} 次请求，耗时 {} ms", atomicInteger.get(), System.currentTimeMillis() - begin);
        return atomicInteger.get();
    }

    /**
     * apache HttpClient 默认的并发请求数，将限制请求的并发度，导致即使使用线程池来并发处理请求，1秒耗时的接口也需要超过1秒的时间
     *   发送 10 次请求，耗时 5018 ms
     * @param count
     * @return
     * @throws InterruptedException
     */
    @GetMapping("wrong")
    public int wrong(@RequestParam(value = "count", defaultValue = "10") int count) throws InterruptedException {
        return sendRequest(count, () -> httpClient1);
    }

    /**
     * 使用重新配置的请求并发数和最大请求并发数，提供处理效率
     *    发送 10 次请求，耗时 1089 ms
     * @param count
     * @return
     * @throws InterruptedException
     */
    @GetMapping("right")
    public int right(@RequestParam(value = "count", defaultValue = "10") int count) throws InterruptedException {
        return sendRequest(count, () -> httpClient2);
    }

    @GetMapping("server")
    public int server() throws InterruptedException {
        TimeUnit.SECONDS.sleep(1);
        return 1;
    }
}
